// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/buffer/buffer.h>
#include <pollux/dwio/common/buffer_util.h>
#include <pollux/dwio/parquet/reader/nested_structure_decoder.h>
#include <pollux/vector/type_aliases.h>

#include <melon/benchmark.h>

using namespace kumo::pollux;
using namespace kumo::pollux::parquet;

class NestedStructureDecoderBenchmark {
 public:
  NestedStructureDecoderBenchmark(uint64_t numValues)
      : numValues_(numValues),
        definitionLevels_(new unsigned char[numValues]()),
        repetitionLevels_(new unsigned char[numValues]()),
        pool_(memory::memoryManager()->addLeafPool()) {}

  void setUp(uint16_t maxDefinition, uint16_t maxRepetition) {
    dwio::common::ensureCapacity<uint64_t>(
        nullsBuffer_, numValues_ / 64 + 1, pool_.get());
    dwio::common::ensureCapacity<vector_size_t>(
        offsetsBuffer_, numValues_ + 1, pool_.get());
    dwio::common::ensureCapacity<vector_size_t>(
        lengthsBuffer_, numValues_, pool_.get());

    populateInputs(maxDefinition, maxRepetition);
  }

  uint64_t numValues_;

  std::unique_ptr<uint8_t> definitionLevels_;
  std::unique_ptr<uint8_t> repetitionLevels_;
  std::shared_ptr<memory::MemoryPool> pool_;

  BufferPtr offsetsBuffer_;
  BufferPtr lengthsBuffer_;
  BufferPtr nullsBuffer_;

 private:
  void populateInputs(uint16_t maxDefinition, uint16_t maxRepetition) {
    auto def = definitionLevels_.get();
    auto rep = repetitionLevels_.get();
    for (int i = 0; i < numValues_; i++) {
      def[i] = rand() % maxDefinition;
      rep[i] = rand() % maxRepetition;
    }
  }
};

BENCHMARK(randomDefs) {
  melon::BenchmarkSuspender suspender;

  auto numValues = 1'000'000;
  auto maxDefinition = 9;
  auto maxRepetition = 4;

  NestedStructureDecoderBenchmark benchmark(numValues);
  benchmark.setUp(maxDefinition, maxRepetition);

  suspender.dismiss();

  int64_t numCollections = NestedStructureDecoder::readOffsetsAndNulls(
      benchmark.definitionLevels_.get(),
      benchmark.repetitionLevels_.get(),
      numValues,
      maxDefinition / 2,
      maxRepetition / 2,
      benchmark.offsetsBuffer_,
      benchmark.lengthsBuffer_,
      benchmark.nullsBuffer_,
      *benchmark.pool_);

  melon::doNotOptimizeAway(numCollections);
}

int main(int /*argc*/, char** /*argv*/) {
  memory::MemoryManager::initialize({});
  melon::runBenchmarks();
  return 0;
}
